/**
 * @class: TokenBuket
 * @description:
 * @author: fs
 * @create: 2024/6/26 13:57
 **/

import java.util.concurrent.atomic.AtomicLong;

public class TokenBucket {
    private final long capacity; // 令牌桶的容量
    private final long refillTokens; // 每次填充的令牌数量
    private final long refillInterval; // 填充间隔（毫秒）
    private final AtomicLong tokens; // 当前令牌数量
    private long lastRefillTime; // 上次填充时间

    public TokenBucket(long capacity, long refillTokens, long refillInterval) {
        this.capacity = capacity;
        this.refillTokens = refillTokens;
        this.refillInterval = refillInterval;
        this.tokens = new AtomicLong(capacity);
        this.lastRefillTime = System.currentTimeMillis();
    }

    // 尝试获取令牌
    public synchronized boolean tryConsume(long numTokens) {
        refill();

        if (tokens.get() >= numTokens) {
            tokens.getAndAdd(-numTokens);
            return true;
        }

        return false;
    }

    // 填充令牌
    private void refill() {
        long now = System.currentTimeMillis();
        long intervalSinceLast = now - lastRefillTime;

        if (intervalSinceLast > refillInterval) {
            long newTokens = (intervalSinceLast / refillInterval) * refillTokens;
            tokens.set(Math.min(capacity, tokens.get() + newTokens));
            lastRefillTime = now;
        }
    }

    public static void main(String[] args) {
        //TokenBucket tokenBucket = new TokenBucket(10, 1, 1000);
        //
        //Runnable task = () -> {
        //    if (tokenBucket.tryConsume(1)) {
        //        System.out.println(Thread.currentThread().getName() + " - Token consumed");
        //    } else {
        //        System.out.println(Thread.currentThread().getName() + " - Not enough tokens");
        //    }
        //};
        //
        //for (int i = 0; i < 20; i++) {
        //    new Thread(task).start();
        //    try {
        //        Thread.sleep(500); // 模拟请求间隔
        //    } catch (InterruptedException e) {
        //        e.printStackTrace();
        //    }
        //}
    }

}

